import torch
import numpy as np
import cv2
from CVAE import CVAE  # 假设你的模型定义在 CVAE.py 文件中
from dml_torch.dml import TripletModel  # 假设你的模型定义在 TripletModel.py 文件中
from transformers import GPT2Tokenizer, GPT2LMHeadModel

# 加载模型
embedding_dim = 128  # 假设 embedding_dim 是 128
latent_dim = 32  # 假设 latent_dim 是 32

cvae_model = CVAE(input_dim=embedding_dim, latent_dim=latent_dim)
cvae_model.load_state_dict(torch.load('src/model_pth/cvae_model.pth', map_location=torch.device('cpu')))
cvae_model.eval()

triplet_model = TripletModel()
triplet_model.load_state_dict(torch.load('src/model_pth/triplet_model_iu_xray.pth', map_location=torch.device('cpu')))
triplet_model.eval()

# 加载 GPT-2 模型和分词器
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
gpt2_model = GPT2LMHeadModel.from_pretrained('gpt2')
gpt2_model.eval()

# 加载图像
image_path = 'data/iu_xray/iu_xray/images/CXR1_1_IM-0001/0.png'  # 替换为你的图像路径
image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
image = cv2.resize(image, (32, 32))  # 调整图像大小
image = image.astype('float32') / 255  # 归一化
image = np.expand_dims(image, axis=-1)  # 增加通道维度
image = np.expand_dims(image, axis=0)  # 增加 batch 维度
image_tensor = torch.tensor(image, dtype=torch.float32)

# 生成图像嵌入
with torch.no_grad():
    image_embedding = triplet_model(image_tensor)

# 生成文本嵌入
with torch.no_grad():
    new_text_embedding = cvae_model.decode(torch.randn(1, latent_dim), image_embedding)

print("Generated Text Embedding:", new_text_embedding)

# 解码文本嵌入
def decode_text_embedding(text_embedding):
    # 将文本嵌入转换为 GPT-2 模型的输入格式
    input_ids = tokenizer.encode(text_embedding, return_tensors='pt')
    
    # 生成文本
    with torch.no_grad():
        outputs = gpt2_model.generate(input_ids, max_length=50, num_return_sequences=1)
    
    # 解码生成的文本
    decoded_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return decoded_text

decoded_text = decode_text_embedding(new_text_embedding)
print("Decoded Text:", decoded_text)